from flask import Flask, send_file, request
import base64
from PIL import Image
import io

from lama_predict import main as lama_predict

import os
import yaml
from omegaconf import OmegaConf

cwd = os.getcwd()
print(cwd)

config_path = os.path.join(cwd, "configs/prediction/default.yaml")
with open(config_path, 'r') as f:
    config = OmegaConf.create(yaml.safe_load(f))

config.model.path = os.path.join(cwd, "big-lama")
config.indir = os.path.join(cwd, "web_server_input")
config.outdir = os.path.join(cwd, "web_server_output")
config.refine = False

app = Flask(__name__)


@app.route("/api/v2/image", methods=["GET", "POST"])
def echo_image():
    # Get the image data from the request body
    json_dict = request.get_json()
    print(type(json_dict))
    # Get the value of the "image" key, which is the base64 encoded image data
    base64_image_data = json_dict["image"]
    # print(base64_image_data[0:500])

    image_bytes = base64.b64decode(base64_image_data)
    image_stream = io.BytesIO(image_bytes)
    image = Image.open(image_stream)
    print(image.format_description)
    if not os.path.exists("web_server_input"):
        os.makedirs("web_server_input")
    image.save("web_server_input/server.png")

    base64_mask_data = json_dict["mask"]
    image_bytes = base64.b64decode(base64_mask_data)
    image_stream = io.BytesIO(image_bytes)
    mask = Image.open(image_stream)
    print(mask.format_description)
    print(mask.format)
    print(mask.size)
    print(mask.mode)
    if mask.mode != "L":
        mask = mask.convert("L")
    if not os.path.exists("web_server_input"):
        os.makedirs("web_server_input")
    mask.save("web_server_input/server_mask.png")

    # Apply the mask to the image
    # Create a new transparent image with the same size and mode as the image
    transparent = Image.new(image.mode, image.size, (0, 0, 0, 0))
    # Composite the image and the transparent image using the mask
    masked_image = Image.composite(image, transparent, mask)
    masked_image.save("server_masked_image.png")

    # Convert the masked image to bytes and create a new stream
    masked_image_stream = io.BytesIO()
    masked_image.save(masked_image_stream, format='PNG')
    masked_image_stream.seek(0)

    lama_predict(config)

    with open("web_server_output/server_mask.png", "rb") as image_file:
        image_bytes = image_file.read()
        image_inpainted_stream = io.BytesIO(image_bytes)
        print(image.format_description)
    image_inpainted_stream.seek(0)

    return send_file(image_inpainted_stream, mimetype="image/png")


if __name__ == "__main__":
    app.run(debug=True, port=9171)
